from sparkai.llm.llm import ChatSparkLLM, ChunkPrintHandler
from sparkai.core.messages import ChatMessage

XUNFEI_APP_ID = '580ffccb'
XUNFEI_API_KEY = 'ed7b5857af7f77e267a3d0ce0b87b0f0'
XUNFEI_SECRET_KEY = 'NmIzODJhN2YwZjFiNDFhNmE4ZjU5YzI3'

class XunfeiChat:
    def __init__(self, app_id = XUNFEI_APP_ID, api_key=XUNFEI_API_KEY, secret_key=XUNFEI_SECRET_KEY):
        self.app_id = app_id
        self.api_key = api_key
        self.secret_key = secret_key
        self.access_token = None

    # 获取Token
    def get_token(self):
        return self.access_token

    # 调用模型
    def do_chat(self, message):
        spark = ChatSparkLLM(
            spark_api_url='wss://spark-api.xf-yun.com/v4.0/chat',
            spark_app_id = self.app_id,
            spark_api_key = self.api_key,
            spark_api_secret = self.secret_key,
            spark_llm_domain = '4.0Ultra',
            streaming = False
        )
        messages = [ChatMessage(
            role = "user",
            content = message,
            streaming = False
        )]
        params = {
            "chat": {
                "show_ref_label": True
            }
        }
        handler = ChunkPrintHandler()
        response = spark.generate([messages], params=params, callbacks=[handler])
        result = response.generations[0][0].text

        return result

if __name__ == '__main__':
    chat_bot = XunfeiChat()
    response = chat_bot.do_chat("你好")
    print(response)  
